# -*- coding: utf-8 -*-
"""
处理导入gephi的数据，主要修改边的权值
"""
import csv

gephi_node = 'gephi_node.csv'
gephi_edge = 'gephi_edge.csv'

dict_node = {}
dict_edge = {}

# 初始化字典
with open(gephi_node) as f1:
    reader = csv.reader(f1)
    # print(list(reader))

    for row in reader:
        # 第一行是表头
        if reader.line_num > 1:
            # 行号从1开始
            # print(reader.line_num, row)

            # 构成字典，键为label
            dict_node[row[1]] = row

with open(gephi_edge) as f2:
    reader = csv.reader(f2)
    for row in reader:
        if reader.line_num > 1:
            # 字典不包含键
            if row[0] not in dict_edge.keys():
                dict_edge[row[0]] = {}
            # 构成两层字典结构
            dict_edge[row[0]][row[1]] = row[2]

    # print(dict_edge)

# 修改边权
# 原边权表示边的两个顶点同时出现的次数
with open(gephi_edge) as f2:
    reader = csv.reader(f2)
    for row in reader:
        if reader.line_num > 1:
            try:
                dict_edge[row[0]][row[1]] = float(row[2]) / \
                    (float(dict_node[row[0]][2]) + float(dict_node[row[1]][2]))
            except ValueError:
                print('invalid input')

# print(dict_edge)

# 将处理好的边值导出

# 使用数字和字符串的数字都可以
label = ['Source', 'Target','Weight']

with open('output_edge.csv', 'w', newline='') as f:
    writer = csv.writer(f)
    writer.writerow(label)
    for source in dict_edge:
        for target, weight in dict_edge[source].items():
            writer.writerow([source, target, weight])
